from tinydb import TinyDB, Query
import os


class VideoInfoDb:
    def __init__(self, db_path):
        self.db = TinyDB(db_path)
        print(f"Database path: {db_path}")
        print(f"Database size: {len(self.db)}")

    def update_tags(self, timestamp, video_path, video_tags: list, bag_path=""):
        items = self.db.search(Query().id == timestamp)
        if len(items) > 0:
            print("Update video tags")
            item = items[0]
            tags = item["video_tags"]
            tags_set = set(tags)
            tags_set.update(video_tags)
            tags = list(tags_set)
            self.db.update(
                {"video_tags": list(set(tags))}, Query().video_path == video_path
            )
            if bag_path != "":
                self.db.update({"bag_path": bag_path}, Query().video_path == video_path)
        else:
            self.insert_video_info(timestamp, video_path, video_tags, bag_path)

    def insert_video_info(self, time_stamp, video_path, video_tags: list, bag_path=""):
        self.db.insert(
            {"id": time_stamp, "video_path": video_path, "video_tags": video_tags, "bag_path": bag_path}
        )

    def search_videos_by_tag(self, tag):
        Video = Query()
        paths = []
        search_results = self.db.search(Video.video_tags.any(tag))
        for result in search_results:
            paths.append(result["video_path"])
        return paths


    def get_tags_count(self, tags: list):
        Video = Query()
        tags_count = []
        for tag in tags:
            count = len(self.db.search(Video.video_tags.any(tag)))
            tags_count.append({tag:count})
        return tags_count

    def search_videos_by_tags(self, tags: list):
        Video = Query()
        paths = []
        all_tags = dict()
        search_results = self.db.search(Video.video_tags.all(tags))
        for result in search_results:
            paths.append(result["video_path"])
            for tag in result["video_tags"]:
                if tag not in all_tags:
                    all_tags[tag] = 0
                all_tags[tag] += 1

        tags = []
        for tag, count in all_tags.items():
            if count != len(paths):
                tags.append(tag)
        return paths, tags

    def get_video_tags(self, video_path):
        Video = Query()
        items = self.db.search(Video.video_path == video_path)
        if len(items) > 0:
            return items[0]["video_tags"]
        return []

    def get_all_tags(self):
        tags = dict()
        for video_info in self.db.all():
            for tag in video_info["video_tags"]:
                if tag not in tags:
                    tags[tag] = 0
                tags[tag] += 1
        all_tags = []
        for tag, count in tags.items():
            if count == len(self.db):
                continue
            all_tags.append(tag)
        return all_tags

    def get_all_video_paths(self):
        paths = []
        for video_info in self.db.all():
            paths.append(video_info["video_path"])
        return paths

    def get_videos_tags(self, video_folder, videos: list):
        tags = set()
        for video in videos:
            video_path = os.path.join(video_folder, video)
            video_tags = self.get_video_tags(video_path)
            tags.update(video_tags)
        return list(tags)

    def get_video_bags(self, video_path):
        Video = Query()
        items = self.db.search(Video.video_path == video_path)
        if len(items) > 0:
            if "bag_path" in items[0]:
                return items[0]["bag_path"]
        return ""

    def delete_by_video_path(self, video_path):
        self.db.remove(Query().video_path == video_path)

    def clear_bd(self, folder):
        import tqdm
        pbar = tqdm.tqdm(total=len(self.db), desc="Clearing DB")
        for video_info in self.db.all():
            if not os.path.exists(video_info["video_path"]):
                self.db.remove(doc_ids=[video_info.doc_id])
            pbar.update(1)
        pbar.close()
    def close(self):
        self.db.close()


if __name__ == "__main__":
    import sys
    db_json = sys.argv[1] if len(sys.argv) > 1 else "video_info.json"
    db = VideoInfoDb(db_json)
    db.clear_bd(os.path.dirname(db_json))
    print("Clear!")
    try:
        while True:
            pass
    except KeyboardInterrupt:
        db.stop()
        
